import sys, pygame, random, datetime, shelve
import numpy as np
from collections import defaultdict
from decimal import Decimal
from pygame.locals import *
import time


'''
Ant with basic move schemas.  Boards with goal spaces.  Contains functions which
can breed ants.  By selecting the ants which succeed basic challenges and breeding
them, they can obtain spaces that take many (many!) more trials if by random sampling.

use demo1() and demo2() like so:
dna, board = demo1()
..wait..
main(30,30, dna[0 or 1 .. or 9], board)

'''
shelf = shelve.open("boards")
numCycles = [200, 600]
boardList1 = [shelf['maze2'],
              shelf['maze1']]

boardList2 = [shelf['maze2'],
              shelf['maze1'],
              shelf['maze3'],
              shelf['maze4']]

boardList3 = [shelf['maze2'],
              shelf['maze2_2'],
              shelf['maze2_3'],
              shelf['maze2_4']]


########
########
#   The base code for the pygame portion of the ant simulation is from another script.  Virtually the only
#   original component are these global variables.  I'm pretty sure I nested them in the functions where
#   they belong, but since I wasn't planning on developing this script much more, I didn't really worry about them.

dirs = (
        (-1, 0),        #Left
        (0,  1),        #Up
        (1,  0),        #Right
        (0, -1)         #Down
        )
antdir = 1
cellSize = 12
rows = 10
cols = 20
#numCells = 20           # length of the side of the board
background = 0, 0, 0    # background colour; black here
foreground = 23, 23, 23 # foreground colour; the grid's colour; dark gray here
textcol = 177, 177, 177 # the colour of the step display in the upper left of the screen
antwalk = 44, 88, 44    # the ant's trail; greenish here
antant = 222, 44, 44    # the ant's colour; red here
fps = 1.0 / 100      # time between steps; 1.0 / 40 means 40 steps per second
cycles = 1000
boardname = "random"

numPher = 2
numSchema = 4
pherYel = 255, 255, 128
pherOra = 255, 128, 0
pherBlu = 128, 255, 255
wall = 128, 64, 0
sugar = 255, 0 , 255
pherList = [background, pherYel, pherBlu, sugar, wall]
########
########



        
def getdna(data):
    dnaList = []
    for i in range(len(data[2])):
        dnaList.append(data[2][i]["dna"])
    return dnaList
        
def gendna(num):
    dnaList = [[(random.randint(1, numPher),                     # 2 is first pher, 
        random.randint(0,3),                            # Direction turned
        random.randint(0,1)) for i in range(10)] for ii in xrange(num)]         # DNA codes (pheromone dropped, spaces moved))
    return dnaList

def getGeneDist(dnaList, distList = None):
    dnaKey = [(kk,
           jj,
           ii) for kk in xrange(1,3) for jj in xrange(4) for ii in xrange(2)] #16 total
    numGene = len(dnaList[0])
    numSchemas = len(dnaKey)
    if not distList: distList = [0 for ii in xrange(numGene*numSchemas)]
    for dna in dnaList:
        geneNum = 0
        for gene in dna:
            distList[geneNum*numSchemas + dnaKey.index(gene)] += 1
            geneNum+=1
    return distList



def getUnique(dnaList, compList = None):
    # not order preserving
    newList = []
    for x in dnaList:
        if x not in newList:
            newList.append(x)
    if compList == True:
        for x in compList:
            if x not in newList:
                newList.append(x)
    return newList


def getUniqueCount(dnaList, retString = False, distList = None):
    d = defaultdict(int)
    for i in dnaList:
        d[str(i)] += 1
    if retString == True: listed = [d.iteritems()]
    else: listed = [(eval(dna[0]), dna[1]) for dna in d.iteritems()]
    return listed

            
def ss(step, board, pherList, antx, anty, fname=None):
    '''capture board, save as fname'''
    pygame.init()
    rows = len(board[0])
    cols = len(board)
    size = width, height = cols * cellSize, rows * cellSize
    pygame.display.set_caption("Langton's Ant")
    screen = pygame.display.set_mode(size) # Screen is now an object representing the window in which we paint
    screen.fill(background)

    for i in xrange(1, cols):
        pygame.draw.line(screen, foreground, (i * cellSize, 1), (i * cellSize, rows * cellSize), 2)
    for i in xrange(1, rows):
        pygame.draw.line(screen, foreground, (1, i * cellSize), (cols * cellSize, i * cellSize), 2)



    for x in xrange(cols):
        for y in xrange(rows):
            screen.fill(pherList[board[x][y]], pygame.Rect(x * cellSize + 1, y * cellSize + 1, cellSize - 2, cellSize - 2))
    if antx:
        screen.fill(antant, pygame.Rect(antx * cellSize + 1, anty * cellSize + 1, cellSize -2, cellSize -2))
    if step:
        font = pygame.font.Font(None, 36)
        text = font.render("%d " % (step), True, textcol, background)
        screen.blit(text, (10, 10))

    if fname:
        pygame.image.save(screen, fname)
    else:
        pygame.display.flip() # IMPORTANT: No changes are displayed until this function gets called
    return


def randSug(board, cols, rows, sugar):
    sideX = random.randint(0,1)
    sideY = random.randint(0,1)
    sugX = random.randint(1, cols-1)*(-1**sideX)
    sugY = random.randint(1, rows-1)*(-1**sideY)
    board[sugX][sugY] = sugar
    return board

def placeSug(board, sugX, sugY, sugar):
    board[sugX][sugY] = sugar


def boardselection(cycles, threshhold, boardList, numBoardUse, liveCrit = None,
           randomize = False, dnaList = None, ants = -1, storeLive = True,
           storeDead = False, reportByRound = False):
    '''Runs each ant through a progression of boards.  Ants go through numBoardUse trials, and those which
succeed on all of them are stored.  liveCrit was intended to select ants which succeed on fewer boards, but
hasn't been implemented yet (or ever?).
'''
    live = 0
    die = 0
    roundReport = [0]*numBoardUse
    liveList = []
    deadList = []
    if dnaList: ants = len(dnaList)
    while live < threshhold and (live + die) != ants:
        if not dnaList:
            dna = [(random.randint(1, numPher),                     # 2 is first pher, 
                random.randint(0,3),                            # Direction turned
                random.randint(0,1)) for i in range(10)]         # DNA codes (pheromone dropped, spaces moved))
        else:
            dna = dnaList[live + die]
        success = 0
        for boardNum in xrange(numBoardUse):
            newboard = [[boardList[boardNum][x][y] for y in xrange(len(boardList[boardNum][0]))] for x in xrange(len(boardList[boardNum]))]

            data = sim(cycles[boardNum],
                   dna,
                   len(newboard[0]),
                   len(newboard),
                   newboard, sug = False, ssWinner = False)


            if data["live"] == False:
                if storeDead == True: deadList.append(data, success)
                die += 1
                if (die + live) % 50000 == 0:
                    print "Dead ==> " + str(die)
                    print "Live ==> " + str(live)
                    time.sleep(1)
                break
            if reportByRound: roundReport[success] += 1
            success += 1

        if success == numBoardUse:
            if storeLive == True: liveList.append(data)
            live += 1
    print live
    print roundReport
    return (live, die, liveList, deadList)
    
def breed(dnaSample, numSwaps, numProgeny, keepParents = True):
    progList = []
    for child in xrange(numProgeny):
        parents = random.sample(dnaSample, 2)
        #inherit = [random.randint(0,1) for xrange(len(parent[0]))]
        new_dna = [ parents[ random.randint(0,1)][gene] for gene in xrange(len(parents[0]))]
        progList.append(new_dna)
    if keepParents == True: return progList.extend(dnaSample)
    else: return progList
            
            

def selection(cycles, threshhold, rows, cols, dnaList=None, ants = -1):
    '''Selects for ants which find sugar under a set number of cycles.  Set ants = n > 0 to run n ants.  Otherwise,
runs until succesful ants = threshold is reached.'''

    live = 0
    die = 0
    deadList = []
    liveList = []
    while live < threshhold:
        dna = [(random.randint(1, numPher),                     # 2 is first pher, 
            random.randint(0,3),                            # Direction turned
            random.randint(0,1)) for i in range(10)]       # DNA codes (pheromone dropped, spaces moved))

        fname = "Screenshots/Alive/%s.jpg" %(live)
        data = sim(cycles, dna, rows, cols, board=None, fname = fname, sug = True)
        if data["live"] == False:
            deadList.append(data)
            die += 1
            if die % 1000 == 0:
                print die
                time.sleep(1)
        else:
            liveList.append(data)
            live += 1
            print live

    return (live, die, liveList, deadList)
    
            

def sim(cycles, dna, rows, cols, board = None, fname=None, sug = False, geneUse = True, ssWinner = True):
    '''Runs through identical simulation as main, except returns a dict of figures:
    dna - ants dna
    boardname - a relic from the past, and I'm afraid to remove it (right now)
    geneAct - 10 slot list with gene activation counts.
    live - bool, True for found sugar
    "sugar coord" - where sugar was located
'''

    boardname = "random"
    size = width, height = cols * cellSize, rows * cellSize

    antx, anty = cols / 2, rows / 2
    antdir = 0
    if board == None:
        board = [[False] * rows for e in xrange(cols)]
        for ii in xrange(cols):
            board[ii][0] = -1
            board[ii][rows-1] = -1
        for ii in xrange(rows):
            board[0][ii] = -1
            board[cols-1][ii] = -1
    sugX = 0
    sugY = 0
    if sug == True:
        sugX = random.randint(1, cols/3)
        sugY = random.randint(1, rows/3)
        board[sugX][sugY] = -2
    
    step = 1
    pause = False
    foundSug = False
    memory = [0]
    
    if geneUse == True: geneActivation = [0 for e in xrange(10)]
    while step <= cycles:
        #Move Instructions
        lookForward = board[antx + dirs[antdir][0]][anty + dirs[antdir][1]]

            
        if lookForward == -2:                           #finds the sugar
            foundSug = True
            if ssWinner == True:
                now = datetime.datetime.today()
                name = now.strftime("%d-%m--%M.%S")
                if not fname: fname = "Screenshots/Alive/%s+%s.jpg" %(name,step)
                ss(step, board, pherList, antx, anty, fname)

            return {"dna" : dna,
                "step": step,
                "boardname" : boardname,
                "geneAct" : geneActivation,
                "live" : foundSug,
                "sugar coord" : (sugX, sugY)}
            

            
        if lookForward == 0:
            gene = lookForward +memory[0]
        elif lookForward == -1:
            gene = -1*memory[0]
        else: gene = lookForward*numPher + memory[0]        #space in front of ant + memory
        


        if geneUse == True: geneActivation[gene] += 1


        board[antx][anty] = dna[gene][0]
        antdir = (antdir + dna[gene][1]) % 4
        memory[0] = board[antx][anty]

        if board[(antx + dirs[antdir][0]*dna[gene][2]) % rows][(anty + dirs[antdir][1]*dna[gene][2]) % cols] == -1:
            antdir = (antdir + dna[-3][1]) % 4                              #activate gene -3 if wall in way of move
            if board[antx + dirs[antdir][0]][anty + dirs[antdir][1]] != -1:
                antx = (antx + dirs[antdir][0]*dna[gene][2]) % cols
                anty = (anty + dirs[antdir][1]*dna[gene][2]) % rows
                
        else:
            antx = (antx + dirs[antdir][0]*dna[gene][2]) % cols
            anty = (anty + dirs[antdir][1]*dna[gene][2]) % rows


        
        #if step % 2500 == 0:
         #   now = datetime.datetime.today()
          #  name = now.strftime("%d-%m-%y--%M.%S")
           # fname = "%s+%s.jpg" %(name,step)
            #ss(step, board, pherList, antx, anty, fname)
            
        step += 1
    return {"dna" : dna,
        "step": step,
        "boardname" : boardname,
        "geneAct" : geneActivation,
        "live" : foundSug,
        "sugar coord" : (sugX, sugY)}
            



def main(rows, cols, dna = None, board = None):
    '''Ant simulation.  Displays ant movements on a board.  Most randomly generated ants aren't
very exciting.
        spacebar -- pause
        some key -- screen shot (tell me if you figure out what event 115 is in pygame)
'''

    pygame.init()

    size = width, height = cols * cellSize, rows * cellSize

    pygame.display.set_caption("Ant Simulation")

    if dna == None:
        dna = [(random.randint(1, numPher),                     # 2 is first pher, 
            random.randint(0,3),                            # Direction turned
            random.randint(0,1)) for i in range(10)]         # DNA codes (pheromone dropped, spaces moved))

        
    screen = pygame.display.set_mode(size) # Screen is now an object representing the window in which we paint
    screen.fill(background)
    pygame.display.flip() # IMPORTANT: No changes are displayed until this function gets called

    for i in xrange(1, cols):
        pygame.draw.line(screen, foreground, (i * cellSize, 1), (i * cellSize, rows * cellSize), 2)
    for i in xrange(1, rows):
        pygame.draw.line(screen, foreground, (1, i * cellSize), (cols * cellSize, i * cellSize), 2)
    pygame.display.flip() # IMPORTANT: No changes are displayed until this function gets called

    font = pygame.font.Font(None, 36)

    antx, anty = cols / 2, rows / 2
    antdir = 0
    if board == None:
        board = [[False] * rows for e in xrange(cols)]
        for ii in xrange(cols):
            board[ii][0] = -1
            board[ii][rows-1] = -1

        for ii in xrange(rows):
            board[0][ii] = -1
            board[cols-1][ii] = -1
    else:
        board = [[board[x][y] for y in xrange(len(board[0]))] for x in xrange(len(board))]
        for x in xrange(cols):
            for y in xrange(rows):
                screen.fill(pherList[board[x][y]], pygame.Rect(x * cellSize + 1, y * cellSize + 1, cellSize - 2, cellSize - 2))
        

    step = 1
    pause = False
    memory = [0]
    while True:
        for event in pygame.event.get():
                if event.type == QUIT:
                    pygame.quit()
                    return
                elif event.type == KEYUP:
                    if event.key == 32: # If space pressed, pause or unpause
                        pause = not pause
                    elif event.key == 115:
                        pygame.image.save(screen, "Step%d.bmp" % (step))
                        now = datetime.datetime.today()
                        name = now.strftime("%d-%m-%y--%M.%S")
                        fname = "%s+%s.bmp" %(name,step)
                        ss(None, board, pherList, antx, anty, fname)

        if pause:
            time.sleep(fps)
            continue

        text = font.render("%d " % (step), True, textcol, background)
        screen.blit(text, (10, 10))


        #Move Instructions
        lookForward = board[antx + dirs[antdir][0]][anty + dirs[antdir][1]]
        if lookForward == -2:         #if ant hits sugar tile during demo,
            lookForward = 0           #treat like ground
            gene = lookForward + memory[0]

            
        if lookForward == 0:
            gene = lookForward +memory[0]
        elif lookForward == -1:
            gene = -1*memory[0]
        else: gene = lookForward*numPher + memory[0]        #space in front of ant + memory


        print gene
        board[antx][anty] = dna[gene][0]                #Lay hormone
        screen.fill(pherList[dna[gene][0]], pygame.Rect(antx * cellSize + 1, anty * cellSize + 1, cellSize - 2, cellSize - 2))
        antdir = (antdir + dna[gene][1]) % 4
        memory[0] = board[antx][anty]

        #if ant tries to walk into wall, apply turn using gene -3
        if board[(antx + dirs[antdir][0]*dna[gene][2]) %rows ][(anty + dirs[antdir][1]*dna[gene][2]) % cols] == -1:
            antdir = (antdir + dna[-3][1]) % 4
            if board[antx + dirs[antdir][0]][anty + dirs[antdir][1]] != -1:
                antx = (antx + dirs[antdir][0]*dna[gene][2]) % cols
                anty = (anty + dirs[antdir][1]*dna[gene][2]) % rows
                print -3
                
        else:
            antx = (antx + dirs[antdir][0]*dna[gene][2]) % cols
            anty = (anty + dirs[antdir][1]*dna[gene][2]) % rows

        # The current square (i.e. the ant) is painted a different colour
        screen.fill(antant, pygame.Rect(antx * cellSize + 1, anty * cellSize + 1, cellSize -2, cellSize -2))

        pygame.display.flip() # No changes are displayed until this function gets called

        step += 1
        time.sleep(fps)

def demo1():
    shelf = shelve.open("boards")
    numCycles = [200, 600]
    boardList1 = [shelf['maze2'],
                  shelf['maze1']]
    data = boardselection(numCycles, 10, boardList1, 2)
    dna = getdna(data)
    return dna, boardList[1]

def demo2():
    shelf = shelve.open("boards")
    boardList1 = [shelf['spikeList'][0]]
    data = boardselection([500], 10, boardList1, 1)
    dna = getdna(data)
    return dna, boardList[0]
    
if __name__ == "__main__":
    main()
